# _*_ encoding: utf-8 _*_
import random
import math
from distance_distribution import load_edges
from assgin_vehicles import assign_vtype
import optparse
from assgin_vehicles import load_vtypes
import os
from generate_start_points import write_format_xml

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def generate_gauss_distribution(avg, std, n):
    '''
    :param avg: the average time of the distribution
    :param std: the stand derivation of the distribution
    :param n: the number of samples
    :return: [sample] * n
    '''

    samples = [random.gauss(avg, std) for i in range(n)]
    return samples


def load_routes(rou_files):
    '''
    :param rou_file: the rou.xml
    :return: the content of rou_file in [ [rou1], [rou2]] format
    '''

    trips = []
    for rou_file in rou_files:
        for events, elem in ET.iterparse(rou_file):
            if events == 'end':
                if elem.tag == "route":
                    edges = elem.attrib['edges']
                    # parse route
                    route = edges.strip().split(" ")
                    trips.append(route)
    return trips


def group_routes(routes, edges, unit=1):
    '''
    :param routes:  [rou_x]
    :param edges:  {edge_id:{shapes:, length:, speed:,}, }
    :param unit:  default is 1KM
    :return:  {group1, group2, ....} distance of groupi is in [unit * i, unit*(i + 1))
    '''
    unit *= 1000
    group = {}
    for i in range(len(routes)):
        accumulate = 0
        for road in routes[i]:
            accumulate += edges[road]['length']
        group_id = int(math.floor(accumulate / unit))
        group.setdefault(group_id, [])
        group[group_id].append(routes[i])
    return group


def test_group_routes():
    routes_file = ["../../data/10221515.rou.xml"]
    routes = load_routes(routes_file)
    net_file = "../../data/hangzhou.net.xml"
    edges = load_edges(net_file)
    groups = group_routes(routes, edges)
    import matplotlib.pyplot as plt
    data = []
    for key in groups:
        data.append((key, len(groups[key])))
    x, y = [], []
    data = sorted(data, key=lambda d: d[0], reverse=False)
    for item in data:
        x.append(item[0])
        y.append(item[1])
    plt.bar(x, y)
    plt.show()


def to_time_string(t):
    return "%02d:%02d" % (t/2, (t%2) * 30)


def test_generate_time_distribution():
    avg = 28800
    std = 3000
    n = 20000
    samples = generate_gauss_distribution(avg, std, n)
    scale = 3600
    x, y = [], []
    pooling = {}
    for i in samples:
        id = int(math.floor(i / scale))
        pooling.setdefault(id, 0)
        pooling[id] += 1

    print('come here!')
    for k, v in sorted(pooling.iteritems(), key=lambda d: d[0], reverse=False):
        x.append(k)
        #x.append(to_time_string(k))
        y.append(v)
    print('over here!')
    import matplotlib.pyplot as plt
    print('x length is {0}'.format(len(x)))
    plt.bar(x, y)
    # xticks = [to_time_string(k) for k in x]
    # plt.xticks(x, xticks)
    plt.tight_layout()
    plt.show()


test_generate_time_distribution()


def generate_od(group, output, vtypes, t_samples, d_avg, d_std):
    root = ET.Element("routes")

    root.attrib['xmlns:xsi'] = "http://www.w3.org/2001/XMLSchema-instance"
    root.attrib['xsi:noNamespaceSchemaLocation'] = "http://sumo.dlr.de/xsd/routes_file.xsd"

    for id, depart in enumerate(t_samples):
        d = int(math.floor( random.gauss(d_avg, d_std) ) )
        l = len(group[d])
        index = random.randint(0, l - 1)
        subnode = ET.SubElement(root, "vehicle")
        attrib = {}
        attrib.setdefault('id', "{0}".format(id))
        attrib.setdefault('depart', "{0}".format(depart))
        attrib.setdefault('departLane', "best")
        attrib.setdefault('departSpeed', '0.00')
        attrib.setdefault('departPos', 'base')
        attrib.setdefault('arrivalPos', 'random')
        attrib.setdefault('type', assign_vtype(vtypes))
        subnode.attrib=attrib
        routenode = ET.SubElement(subnode, "route")
        routenode.attrib.setdefault('edges', " ".join(group[d][index]))
    # tree = ET.ElementTree(root)
    # tree.write(output, encoding="UTF-8", xml_declaration=True)
    write_format_xml(ET.tostring(root, "utf-8"), output)


def get_options(args=None):
    optparser = optparse.OptionParser()

    optparser.add_option("-n", "--net-file", dest='net_file', help="define the road network")
    optparser.add_option("-r", "--route-directory", dest='route_dir', help="the route file directory")
    optparser.add_option("-v", "--vtypes", dest="vtypes_file", help="the vehicle type file")
    optparser.add_option("-o", "--output", dest="output_file", help="the output file")
    optparser.add_option("-a", "--amount", dest="amount", type="int", help="an integer describes the amount of vehicle")
    optparser.add_option("--time-avg", dest="tavg", type="float", help="an float number describes the average /"
                                                                       "value of depart time")
    optparser.add_option("--time-std", dest="tstd", type='float', help="an float number describes the standard /"
                                                                       "derivation of depart time")
    optparser.add_option("--distance-avg", dest="davg", type='float', help="an float number describes the average /"
                                                                           "value of trip distance")
    optparser.add_option("--distance-std", dest="dstd", type='float', help='an float number describes the standard /'
                                                                           'derivation of trip distance')
    optparser.add_option('-u', "--unit", dest="unit", type='int', help="an integer describes the unit of group /"
                                                                       "default is 1KM")

    (options, args) = optparser.parse_args(args=args)
    if not options.net_file or not options.route_dir or not options.vtypes_file or not options.output_file:
        optparser.print_help()
        exit(0)

    if not options.amount or not options.tavg or not options.tstd or not options.davg or not options.dstd:
        optparser.print_help()
        exit(0)

    return options


def main(options=None):

    edges = load_edges(options.net_file)
    root_dirs = os.walk(options.route_dir)
    route_files = []
    for root, dirs, files in root_dirs:
        for f in files:
            if f.endswith(".rou.xml"):
                route_files.append(os.path.join(root, f))

    routes = load_routes(route_files)
    vtypes = load_vtypes(options.vtypes_file)

    samples = generate_gauss_distribution(avg=options.tavg, std=options.tstd, n=options.amount)
    samples.sort()
    for i, t in enumerate(samples):
        samples[i] = int(math.floor(t))
    unit = 1
    if options.unit:
        unit = options.unit
    group = group_routes(routes, edges, unit=unit)
    generate_od(group, options.output_file, vtypes, samples, options.davg, options.dstd)
    print('finished!\n')


def test_dir():
    root_dirs = os.walk("../../data/1030")
    routes_files = []
    for root, dirs, files in root_dirs:
        for f in files:
            if f.endswith(".rou.xml"):
                print("%s" % (os.path.join(root, f)))
                routes_files.append(os.path.join(root, f))

    with open(routes_files[-1], "rb") as f:
        for line in f:
            print line


if __name__ == '__main__':
    main(get_options())
